{-# LANGUAGE AllowAmbiguousTypes #-}

module Q.IR.Interpret
  ( interpret
  , M
  ) where

import Data.Function (on)
import Data.Map.Strict (Map)

import Q.IR.Common
import Q.IR.Core

import qualified Data.Map.Strict as Map

--------------------------------------------------------------------------------
-- Interpreter environment

type Env v = Map Id v

data Context = Context
  { cxVars   :: Env Value
  , cxLabels :: Env LabelDef }

instance Monoid Context where
  mempty = Context
    { cxVars   = Map.empty
    , cxLabels = Map.empty }

instance Semigroup Context where
  cx1 <> cx2 = Context
    { cxVars   = cxVars cx2 `Map.union` cxVars cx1
    , cxLabels = cxLabels cx2 `Map.union` cxLabels cx1 }

class ContextContains v where
  whatsItCalled :: String
  getEnv  :: Context -> Env v
  fromEnv :: Env v -> Context

  extend :: Context -> [(Id, v)] -> Context
  extend cx l = cx <> fromEnv (Map.fromList l)

  find :: Id -> Context -> v
  find x cx = case Map.lookup x (getEnv cx) of
    Just y  -> y
    Nothing -> error $ "unbound " ++ whatsItCalled @v ++ ": " ++ show x

instance ContextContains Value where
  whatsItCalled = "variable"
  getEnv        = cxVars
  fromEnv e     = mempty { cxVars = e }

instance ContextContains LabelDef where
  whatsItCalled = "label"
  getEnv        = cxLabels
  fromEnv e     = mempty { cxLabels = e }

--------------------------------------------------------------------------------
-- Interpreter

type M = IO

interpret :: Prog -> M Value
interpret p0 = evalProg mempty p0

evalProg :: Context -> Prog -> M Value
evalProg cx (PJump k arg) = do
  argVal <- evalExpr cx arg
  case k of
    Return  -> pure argVal
    Label l -> let LabelDef argId _ body = find l cx in
                 evalProg (extend cx [(argId, argVal)]) body

evalProg cx (PDefn defs body) =
  evalProg (extend cx defs) body

evalProg cx (PCases e alts) = do
  n <- valueToInt <$> evalExpr cx e
  case lookup n alts of
    Just p  -> evalProg cx p
    Nothing -> error $ "no matching case for " ++ show n

evalExpr :: Context -> Expr -> M Value
evalExpr cx e0 = return $! eval e0
  where
    eval (EInt i)        = VInt i
    eval (ERef x path)   = follow (find x cx) path
    eval (ECons e1 e2)   = VCons (eval e1) (eval e2)
    eval (EOp2 op e1 e2) = VInt (evalOp op (eval e1) (eval e2))

    evalOp op = applyOp2 op `on` valueToInt
